mport argparse
import json

import pandas as pd


system_prompt = """
我是一个帮你优化日常对话的小助手。

我的主要任务是帮你把话说得更自然、更有吸引力，但还是保持你想表达的意思
"""

user_prompt = """
我在跟我的{target}聊天，我想给对方表达下面这段话，但是不知道怎么说才好。```
{user_input}
```
麻烦帮我按照{tags}，{preferenced_style}的风格改写一下呗。
直接输出给我结果就行了，不需要解释。
"""

template = {"param": {"max_tokens": 2048}, "id": "h_001", "query": ""}

parser = argparse.ArgumentParser()
parser.add_argument(
    "-n",
    "--name",
    type=str,
)
args = parser.parse_args()

df = pd.read_excel(".vscode/test.xlsx")

f = open(".vscode/test.jsonl", "w")
tasks = []
for i, row in df.iterrows():
    user_input = row["input"]
    target = row["target"]
    preferenced_style = row["tags"]
    tags = row["特征"].replace("，", ",").split(",")

    for tag in tags:
        up = user_prompt.format(
            user_input=user_input,
            target=target,
            preferenced_style=preferenced_style,
            tags=tag,
        )
        query = system_prompt + "\n\n" + up

        line = template.copy()
        line["query"] = query
        f.write(json.dumps(line, ensure_ascii=False) + "\n")

f.close()
